#include <ros/ros.h>
#include <pcl/io/pcd_io.h>
#include <pcl/point_types.h>
#include <sensor_msgs/PointCloud2.h>
#include <pcl_conversions/pcl_conversions.h>
#include <rosbag/bag.h>
#include <rosbag/view.h>
#include <boost/foreach.hpp>
#include <sys/stat.h>

#define foreach BOOST_FOREACH

int main(int argc, char** argv)
{
    ros::init(argc, argv, "pcd_extractor");
    ROS_INFO("Program started.");

    if (argc < 3)
    {
        ROS_INFO("Usage: %s <bag_file> <output_dir>", argv[0]);
        return 1;
    }

    std::string bag_file = argv[1];
    std::string output_dir = argv[2];

    ROS_INFO("Bag file: %s", bag_file.c_str());
    ROS_INFO("Output directory: %s", output_dir.c_str());

    rosbag::Bag bag;
    try
    {
        bag.open(bag_file, rosbag::bagmode::Read);
        ROS_INFO("Successfully opened the bag file.");
    }
    catch (rosbag::BagIOException& e)
    {
        ROS_ERROR_STREAM("Failed to open bag file: " << e.what());
        return 1;
    }

    // 检查并创建输出目录
    if (access(output_dir.c_str(), F_OK) == -1) {
        if (mkdir(output_dir.c_str(), 0755) == -1) {
            ROS_ERROR_STREAM("Failed to create output directory: " << output_dir);
            return 1;
        } else {
            ROS_INFO_STREAM("Created output directory: " << output_dir);
        }
    }

    std::vector<std::string> topics{"/livox/lidar"};
    rosbag::View view(bag, rosbag::TopicQuery(topics));
    ROS_INFO("Starting to process messages.");

    int i = 0;
    foreach(rosbag::MessageInstance const m, view)
    {
        sensor_msgs::PointCloud2::ConstPtr cloud_msg = m.instantiate<sensor_msgs::PointCloud2>();
        if (cloud_msg != nullptr)
        {
            ROS_INFO("Processing frame %d.", i);

            pcl::PCLPointCloud2 pcl_pc2;
            pcl_conversions::toPCL(*cloud_msg, pcl_pc2);

            pcl::PointCloud<pcl::PointXYZ>::Ptr cloud(new pcl::PointCloud<pcl::PointXYZ>);
            pcl::fromPCLPointCloud2(pcl_pc2, *cloud);

            ROS_INFO("Converted frame %d to PCL format.", i);

            std::stringstream ss;
            ss << output_dir << "/frame_" << i << ".pcd";
            ROS_INFO("Attempting to save PCD file to: %s", ss.str().c_str());
            if (pcl::io::savePCDFileASCII(ss.str(), *cloud) == 0)
            {
                ROS_INFO("Successfully saved frame_%d.pcd", i);
            }
            else
            {
                ROS_ERROR("Failed to save frame_%d.pcd", i);
            }

            ++i;
        }
        else
        {
            ROS_WARN("Failed to instantiate PointCloud2 message.");
        }
    }

    bag.close();
    return 0;
}